package com.shujia.spark.mllib

import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}

object Demo3ImageMakeData {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[7]")
      .appName("image")
      .config("spark.sql.shuffle.partitions", 8)
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._

    //读取图片数据
    val imageDF: DataFrame = spark
      .read
      .format("image")
      .load("E:\\课件\\机器学习数据\\手写数字\\train")

    imageDF.printSchema()

    //将图片数据转换成向量
    val dataToVector: UserDefinedFunction = udf((data: Array[Byte]) => {
      //对原始的数据做处理
      val array: Array[Double] = data
        .map(x => x.toInt)
        .map(x => {
          if (x >= 0) {
            //黑色像素点
            0.0
          } else {
            //白色像素点
            1.0
          }
        })
      //转换成向量返回,
      //转换成稀疏向量帆布i，减少空间占用
      Vectors.dense(array).toSparse
    })

    //取出图片名称
    val originToName: UserDefinedFunction = udf((origin: String) => {
      origin.split("/").last
    })

    val featuresDF: DataFrame = imageDF
      //取出图片的路径和图片数据
      .select($"image.origin" as "origin", $"image.data" as "data")
      //将原始图片的数据，转换成特征向量
      .select(originToName($"origin") as "name", dataToVector($"data") as "features")


    //读取图片名和标记数据
    val labelDF: DataFrame = spark.read
      .format("csv")
      .option("sep", " ")
      .schema("name STRING, label DOUBLE")
      .load("E:\\课件\\机器学习数据\\手写数字\\train.txt")

    //关联获取目标值和特征向量
    val trainDF: DataFrame = featuresDF
      .join(labelDF.hint("broadcast"), "name")
      .select($"label", $"features")


    //保存数据
    trainDF
      .write
      .format("libsvm") //svm是专门用于机器学习的一个数据格式
      .mode(SaveMode.Overwrite)
      .save("data/image_data")

  }
}
